clear
%% Problem
%   'EqualMinima','Himmelblau','Sixhump','ModifiedRastrigin','Vincent',
%   'Shubert','Composition','IncreasingMinima','Rastrigin','Schaffer'.
ProblemSet.FuncType = 'Himmelblau';
ProblemSet.dimension = 2;
% ProblemSet.k = [3,3,3];
% ProblemSet.BaseFunc = 3;

[d, x_bound, ObjFunc, OptSet, precision_init, precision, ~, MaximalBudget, FileName]...
    = problem_setting(ProblemSet);

rep = 1;
MaximalBudget = 50000;
iepsilon = 4;
epsilon = [0.1,0.01,0.001,0.0001];
%% problem
NewRegionNum = 2;
MaxPartitionDepth_d = ceil(log((x_bound(2,:)-x_bound(1,:))./precision(iepsilon))...
    ./log(NewRegionNum));
MaxPartitionDepth = sum(MaxPartitionDepth_d)
ClearRadius0 = (x_bound(2,:)-x_bound(1,:))./(NewRegionNum.^(MaxPartitionDepth_d));
ClearRadius = 2*min(ClearRadius0); % the maximum distance within a non-partitionable region

Problem.Dimension = d;
Problem.Domain = reshape(x_bound,[1,d*2]);
Problem.Partition = @(Region,SampleSet)Partition_BoxConstraint(Region, NewRegionNum, MaxPartitionDepth, SampleSet);
Problem.Sampling = @(n, region)Sampling_BoxConstraint(n, region, ObjFunc);

%% optimization
% Algorithm parameters
AlgorithmM.n0 = 6;
AlgorithmM.QuantileLevel = 0.2;
AlgorithmM.NewBudget = 3;
AlgorithmM.MaxSampleSize = 15;
AlgorithmM.radius = ClearRadius;
AlgorithmM.StopCriteria = [1,MaximalBudget];
% AlgorithmM.StopCriteria = [2,size(OptSet.sol,1)];
% AlgorithmM.StopCriteria = [3,5000];

%% STAGE ONE PRS-MMO
time_start = tic;
for i=1:rep
[optima, RegionSet_Stage1, SampleSet_Stage1, RegionSampleId_Stage1, FlatRegion] ...
    = prsmmo( Problem, AlgorithmM );
Optima_Stage1 = optima(:,2:end);
end
toc(time_start)
TotalBudget_Stage1 = size(SampleSet_Stage1,1)
%% STAGE TWO LOCAL CONVERGENCE
%% local search
% step = ClearRadius;
% step_threshold = precision(iepsilon)/2;
% [Optima_Stage2, path_Stage2, TotalBudget_Stage2] = LocalSearch_BoxConstraint(ObjFunc,x_bound,Optima_Stage1,step,step_threshold);
% sum(TotalBudget_Stage2)

%% save data
% save(FileName)

%% FIGURES
if rep == 1
    if d==1
        x=x_bound(1,1):(x_bound(2,1)-x_bound(1,1))/1000:x_bound(2,1);
        z=ObjFunc(x');
        z_max=max(z);
        z_min=min(z);
        figure()
        hold on
        plot(x,z)
        for i=1:length(RegionSet_Stage1)
            plot([RegionSet_Stage1(i,3), RegionSet_Stage1(i,3)], [z_min, z_max],'b')
            plot([RegionSet_Stage1(i,4), RegionSet_Stage1(i,4)], [z_min, z_max],'b')
        end
        plot(SampleSet_Stage1(:,2),SampleSet_Stage1(:,1),'.');
        hold off
    end
    if d == 2
        %% Stage one figures
        figure()
        hold on
        rectangle('Position',[Problem.Domain([1,3]),Problem.Domain([2,4])-Problem.Domain([1,3])])
        scatter(SampleSet_Stage1(:,2),SampleSet_Stage1(:,3),2,SampleSet_Stage1(:,1),'filled');
        for i=1:length(RegionSet_Stage1)
            rectangle('Position',[RegionSet_Stage1(i,[3,5]),RegionSet_Stage1(i,[4,6])-RegionSet_Stage1(i,[3,5])])
        end
        title('','Interpreter','latex','String',['$N=',num2str(size(SampleSet_Stage1,1)),'$'])
        xlabel('','Interpreter','latex','String','$x_1$')
        ylabel('','Interpreter','latex','String','$x_2$')
        set(gca,'Fontname','Times New Roman','FontSize',16);
        for i=1:size(Optima_Stage1,1)
            plot(Optima_Stage1(i,2),Optima_Stage1(i,3),'x','MarkerSize',6,'LineWidth',1.5,'MarkerEdgeColor' ,[0.86,0.34,0.07]);
        end
        hold off
        
        %% flat regions
        figure()
        hold on
        rectangle('Position',[Problem.Domain([1,3]),Problem.Domain([2,4])-Problem.Domain([1,3])])
        for i=1:length(FlatRegion)
            rectangle('Position',[RegionSet_Stage1(FlatRegion(i),[3,5]),RegionSet_Stage1(FlatRegion(i),[4,6])-RegionSet_Stage1(FlatRegion(i),[3,5])])
        end
        hold off
        %% Stage two figures
%         figure()
%         hold on
%         rectangle('Position',[Problem.Domain([1,3]),Problem.Domain([2,4])-Problem.Domain([1,3])])
%         plot(Optima_Stage2(:,2),Optima_Stage2(:,3),'*')
%         hold off
%         xlabel('','Interpreter','latex','String','$x_1$')
%         ylabel('','Interpreter','latex','String','$x_2$')
%         title('Local search')
%         set(gca,'Fontname','Times New Roman','FontSize',16);
    end
    if d==3
        %% Stage one figures
        figure()
        scatter3(Optima_Stage1(:,2),Optima_Stage1(:,3),Optima_Stage1(:,4),2,Optima_Stage1(:,1),'filled');
        xlabel('','Interpreter','latex','String','$x_1$')
        ylabel('','Interpreter','latex','String','$x_2$')
        title('PRS-MMO')
        set(gca,'Fontname','Times New Roman','FontSize',16);
        %% Stage two figures
%         figure()
%         scatter3(Optima_Stage2(:,2),Optima_Stage2(:,3),Optima_Stage2(:,4),2,Optima_Stage2(:,1),'filled');
%         xlabel('','Interpreter','latex','String','$x_1$')
%         ylabel('','Interpreter','latex','String','$x_2$')
%         title('Local search')
%         set(gca,'Fontname','Times New Roman','FontSize',16);
    end
end